/**
* @file device.cpp
*
* Copyright (c) Huawei Technologies Co., Ltd. 2019-2020. All Rights reserved.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
*/

#include "acl/acl_rt.h"

#include <cstring>
#include "runtime/dev.h"
#include "runtime/context.h"
#include "runtime/config.h"
#include "runtime/kernel.h"
#include "framework/executor/ge_executor.h"
#include "platform/platform_info.h"
#include "register/stream_manage_func_registry.h"

#include "model/acl_resource_manager.h"
#include "model/model_desc_internal.h"
#include "log_inner.h"
#include "error_codes_inner.h"
#include "toolchain/profiling_manager.h"
#include "toolchain/resource_statistics.h"
#include "set_device_vxx.h"
#include "common_inner.h"

namespace {
    constexpr int32_t DEVICE_UTILIZATION_NOT_SUPPORT = -1;
    constexpr int32_t MODULE_TYPE_VECTOR_CORE = 7;
    constexpr int32_t MODULE_TYPE_AICORE = 4;
    constexpr int32_t INFO_TYPE_CORE_NUM = 3;
    std::mutex g_platformInfoInitMutex;
    std::unordered_set<int32_t> g_platformInfoInitSet;

    int32_t GetAllUtilizations(const int32_t deviceId, const rtTypeUtil_t utilType)
    {
        uint8_t utilRate = 0U;
        const rtError_t rtErr = rtGetAllUtilizations(deviceId, utilType, &utilRate);
        if (rtErr == ACL_ERROR_RT_FEATURE_NOT_SUPPORT) {
            ACL_LOG_WARN("rtGetAllUtilizations not to support this query, utilType = %d, runtime result = %d.",
                         static_cast<int32_t>(utilType), static_cast<int32_t>(rtErr));
            return DEVICE_UTILIZATION_NOT_SUPPORT;
        }
        if (rtErr != RT_ERROR_NONE) {
            ACL_LOG_CALL_ERROR("rtGetAllUtilizations failed, utilType = %d, runtime result = %d.",
                               static_cast<int32_t>(utilType), static_cast<int32_t>(rtErr));
            return DEVICE_UTILIZATION_NOT_SUPPORT;
        }
        ACL_LOG_INFO("successfully execute rtGetAllUtilizations, utilType = %d, utilRate = %u.",
                     static_cast<int32_t>(utilType), utilRate);
        return static_cast<int32_t>(utilRate);
    }
}

namespace acl {
    aclError UpdatePlatformInfoWithDevice(int32_t deviceId)
    {
#ifdef __GNUC__
        // init platform info
        const aclError ret = acl::InitSocVersion();
        if (ret != ACL_SUCCESS) {
            return ret;
        }
        const string &socVersion = acl::GetSocVersion();
        if (fe::PlatformInfoManager::GeInstance().InitRuntimePlatformInfos(socVersion) != 0U) {
            ACL_LOG_INNER_ERROR("[Init][PlatformInfo]init runtime platform info failed, SocVersion = %s",
                                socVersion.c_str());
            return ACL_ERROR_INTERNAL_ERROR;
        }

        const std::unique_lock<std::mutex> lk(g_platformInfoInitMutex);
        if (g_platformInfoInitSet.count(deviceId) > 0U) {
            return ACL_SUCCESS;
        }

        uint32_t aicCnt = 0U;
        auto rtErr = rtGetAiCoreCount(&aicCnt);
        if (rtErr != RT_ERROR_NONE) {
            ACL_LOG_CALL_ERROR("get aicore count failed, runtime result = %d", static_cast<int32_t>(rtErr));
            return ACL_GET_ERRCODE_RTS(rtErr);
        }

        int64_t vecCoreCnt = 0U;
        // some chips has no vector core
        rtErr = rtGetDeviceInfo(static_cast<uint32_t>(deviceId), MODULE_TYPE_VECTOR_CORE,
                                INFO_TYPE_CORE_NUM, &vecCoreCnt);
        if (rtErr != RT_ERROR_NONE) {
            ACL_LOG_CALL_ERROR("get vector core count failed, runtime result = %d", static_cast<int32_t>(rtErr));
            return ACL_GET_ERRCODE_RTS(rtErr);
        }

        int64_t cubeCoreCnt = 0U;
        rtErr = rtGetDeviceInfo(static_cast<uint32_t>(deviceId), MODULE_TYPE_AICORE,
                                INFO_TYPE_CUBE_NUM, &cubeCoreCnt);
        if (rtErr != RT_ERROR_NONE) {
            ACL_LOG_CALL_ERROR("get cube core count failed, runtime result = %d", static_cast<int32_t>(rtErr));
            return ACL_GET_ERRCODE_RTS(rtErr);
        }

        fe::PlatFormInfos platformInfos;
        uint32_t platformRet =
            fe::PlatformInfoManager::GeInstance().GetRuntimePlatformInfosByDevice(deviceId, platformInfos);
        if (platformRet != 0U) {
            ACL_LOG_CALL_ERROR("get runtime platformInfos by device failed, deviceId = %d", deviceId);
            return ACL_ERROR_INTERNAL_ERROR;
        }

        const std::string socInfoKey = "SoCInfo";
        const std::string aicCntKey = "ai_core_cnt";
        const std::string vecCoreCntKey = "vector_core_cnt";
        const std::string cubeCoreCntKey = "cube_core_cnt";
        std::map<std::string, std::string> res;
        if (!platformInfos.GetPlatformResWithLock(socInfoKey, res)) {
            ACL_LOG_CALL_ERROR("unable to get platform result");
            return ACL_ERROR_INTERNAL_ERROR;
        }

        res[aicCntKey] = std::to_string(aicCnt);
        res[vecCoreCntKey] = std::to_string(vecCoreCnt);
        res[cubeCoreCntKey] = std::to_string(cubeCoreCnt);
        platformInfos.SetPlatformResWithLock(socInfoKey, res);
        platformRet =
            fe::PlatformInfoManager::GeInstance().UpdateRuntimePlatformInfosByDevice(deviceId, platformInfos);
        if (platformRet != 0U) {
            ACL_LOG_CALL_ERROR("update runtime platformInfos by device failed, deviceId = %d", deviceId);
            return ACL_ERROR_INTERNAL_ERROR;
        }
        (void)g_platformInfoInitSet.emplace(deviceId);
        ACL_LOG_INFO("Successfully to UpdatePlatformInfoWithDevice, deviceId = %d, aicCnt = %u, vecCoreCnt = %ld, "
                     "cubeCoreCnt = %ld", deviceId, aicCnt, vecCoreCnt, cubeCoreCnt);
#endif
        return ACL_SUCCESS;
    }
}

aclError aclrtSetDevice(int32_t deviceId)
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtSetDevice);
    ACL_ADD_APPLY_TOTAL_COUNT(acl::ACL_STATISTICS_SET_RESET_DEVICE);
    ACL_LOG_INFO("start to execute aclrtSetDevice, deviceId = %d.", deviceId);
    const rtError_t rtErr = rtSetDevice(deviceId);
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("open device %d failed, runtime result = %d.", deviceId, static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    // update platform info
    const auto aclErr = acl::UpdatePlatformInfoWithDevice(deviceId);
    if (aclErr != ACL_SUCCESS) {
        ACL_LOG_WARN("update platform info with device failed, deviceId = %d", deviceId);
    }
    ACL_LOG_INFO("successfully execute aclrtSetDevice, deviceId = %d", deviceId);
    ACL_ADD_APPLY_SUCCESS_COUNT(acl::ACL_STATISTICS_SET_RESET_DEVICE);
    return ACL_SUCCESS;
}

aclError aclrtSetDeviceWithoutTsdVXX(int32_t deviceId)
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtSetDeviceWithoutTsdVXX);
    ACL_ADD_APPLY_TOTAL_COUNT(acl::ACL_STATISTICS_SET_RESET_DEVICE);
    ACL_LOG_INFO("start to execute aclrtSetDeviceWithoutTsdVXX, deviceId = %d.", deviceId);
    const std::string &socVersion = acl::GetSocVersion();
    if (strncmp(socVersion.c_str(), "Ascend910", (sizeof("Ascend910") - 1UL)) != 0) {
        ACL_LOG_INFO("The soc version is not Ascend910, not support");
        return ACL_ERROR_API_NOT_SUPPORT;
    }
    const rtError_t rtErr = rtSetDeviceWithoutTsd(deviceId);
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("open device %d failed, runtime result = %d.", deviceId, static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_INFO("open device %d successfully.", deviceId);
    ACL_ADD_APPLY_SUCCESS_COUNT(acl::ACL_STATISTICS_SET_RESET_DEVICE);
    return ACL_SUCCESS;
}

aclError aclrtResetDevice(int32_t deviceId)
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtResetDevice);
    ACL_ADD_RELEASE_TOTAL_COUNT(acl::ACL_STATISTICS_SET_RESET_DEVICE);
    ACL_LOG_INFO("start to execute aclrtResetDevice, deviceId = %d.", deviceId);
    const rtError_t rtErr = rtDeviceReset(deviceId);
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("reset device %d failed, runtime result = %d.", deviceId, static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_INFO("successfully execute aclrtResetDevice, reset device %d.", deviceId);
    ACL_ADD_RELEASE_SUCCESS_COUNT(acl::ACL_STATISTICS_SET_RESET_DEVICE);
    return ACL_SUCCESS;
}

aclError aclrtResetDeviceForce(int32_t deviceId)
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtResetDeviceForce);
    ACL_ADD_RELEASE_TOTAL_COUNT(acl::ACL_STATISTICS_SET_RESET_DEVICE);
    ACL_LOG_INFO("start to execute aclrtResetDeviceForce, deviceId = %d.", deviceId);
    const rtError_t rtErr = rtDeviceResetForce(deviceId);
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("force reset device %d failed, runtime result = %d.", deviceId, static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_INFO("successfully execute aclrtResetDeviceForce, reset device %d.", deviceId);
    ACL_ADD_RELEASE_SUCCESS_COUNT(acl::ACL_STATISTICS_SET_RESET_DEVICE);
    return ACL_SUCCESS;
}

aclError aclrtResetDeviceWithoutTsdVXX(int32_t deviceId)
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtResetDeviceWithoutTsdVXX);
    ACL_ADD_RELEASE_TOTAL_COUNT(acl::ACL_STATISTICS_SET_RESET_DEVICE);
    ACL_LOG_INFO("start to execute aclrtResetDeviceWithoutTsdVXX, deviceId = %d.", deviceId);
    const std::string &socVersion = acl::GetSocVersion();
    if (strncmp(socVersion.c_str(), "Ascend910", (sizeof("Ascend910") - 1UL)) != 0) {
        ACL_LOG_INNER_ERROR("The soc version is not Ascend910, not support");
        return ACL_ERROR_API_NOT_SUPPORT;
    }
    const rtError_t rtErr = rtDeviceResetWithoutTsd(deviceId);
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("reset device %d failed, runtime result = %d.", deviceId, static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_INFO("successfully execute aclrtResetDeviceWithoutTsdVXX, reset device %d", deviceId);
    ACL_ADD_RELEASE_SUCCESS_COUNT(acl::ACL_STATISTICS_SET_RESET_DEVICE);
    return ACL_SUCCESS;
}

aclError aclrtGetDevice(int32_t *deviceId)
{
    ACL_LOG_INFO("start to execute aclrtGetDevice");
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(deviceId);
    const rtError_t rtErr = rtGetDevice(deviceId);
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_INFO("can not get device id, runtime result = %d.", static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_DEBUG("successfully execute aclrtGetDevice, get device id is %d.", *deviceId);
    return ACL_SUCCESS;
}

aclError aclrtGetRunMode(aclrtRunMode *runMode)
{
    ACL_LOG_INFO("start to execute aclrtGetRunMode");
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(runMode);
    rtRunMode rtMode;
    const rtError_t rtErr = rtGetRunMode(&rtMode);
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("get runMode failed, runtime result = %d.", static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    if (rtMode == RT_RUN_MODE_OFFLINE) {
        *runMode = ACL_DEVICE;
        return ACL_SUCCESS;
    }
    *runMode = ACL_HOST;
    ACL_LOG_INFO("successfully execute aclrtGetRunMode, current runMode is %d.", static_cast<int32_t>(*runMode));
    return ACL_SUCCESS;
}

aclError aclrtSynchronizeDevice()
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtSynchronizeDevice);
    ACL_LOG_INFO("start to execute aclrtSynchronizeDevice");
    const rtError_t rtErr = rtDeviceSynchronize();
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("wait for compute device to finish failed, runtime result = %d.",
            static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_INFO("device synchronize successfully.");
    return ACL_SUCCESS;
}

aclError aclrtSynchronizeDeviceWithTimeout(int32_t timeout)
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtSynchronizeDeviceWithTimeout);
    ACL_LOG_INFO("start to execute aclrtSynchronizeDeviceWithTimeout, timeout %dms", timeout);
    constexpr int32_t defaultTimeout = -1;
    if (timeout < defaultTimeout) {
        ACL_LOG_CALL_ERROR("the timeout of synchronize device is invalid, timeout is %d", timeout);
        return ACL_ERROR_RT_PARAM_INVALID;
    }

    const rtError_t rtErr = rtDeviceSynchronizeWithTimeout(timeout);
    if (rtErr == ACL_ERROR_RT_STREAM_SYNC_TIMEOUT) {
        ACL_LOG_CALL_ERROR("synchronize device timeout, timeout = %dms", timeout);
        return ACL_ERROR_RT_STREAM_SYNC_TIMEOUT;
    } else if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("wait for compute device to finish failed, runtime result = %d.",
            static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_INFO("device synchronize with timeout %dms successfully.", timeout);
    return ACL_SUCCESS;
}

aclError aclrtSetTsDevice(aclrtTsId tsId)
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtSetTsDevice);
    ACL_LOG_INFO("start to execute aclrtSetTsDevice, tsId = %d.", static_cast<int32_t>(tsId));
    if ((tsId != ACL_TS_ID_AICORE) && (tsId != ACL_TS_ID_AIVECTOR)) {
        ACL_LOG_INNER_ERROR("invalid tsId, tsID is %d.", static_cast<int32_t>(tsId));
        return ACL_ERROR_INVALID_PARAM;
    }
    const rtError_t rtErr = rtSetTSDevice(static_cast<uint32_t>(tsId));
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("set device ts %d failed, runtime result = %d.", static_cast<int32_t>(tsId), rtErr);
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_INFO("successfully execute aclrtSetTsDevice, set device ts %d", static_cast<int32_t>(tsId));
    return ACL_SUCCESS;
}

aclError aclrtGetDeviceUtilizationRate(int32_t deviceId, aclrtUtilizationInfo *utilizationInfo)
{
    ACL_LOG_INFO("start to execute aclrtGetDeviceUtilizationRate, device is %d.", deviceId);
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(utilizationInfo);
    aclrtUtilizationExtendInfo *utilizationExtend = utilizationInfo->utilizationExtend;
    if (utilizationExtend != nullptr) {
        ACL_LOG_ERROR("utilizationExtend is reserved parameter, current version needs to be null!");
        return ACL_ERROR_INVALID_PARAM;
    }
    utilizationInfo->cubeUtilization = GetAllUtilizations(deviceId, RT_UTIL_TYPE_AICORE);
    utilizationInfo->vectorUtilization = GetAllUtilizations(deviceId, RT_UTIL_TYPE_AIVECTOR);
    utilizationInfo->aicpuUtilization = GetAllUtilizations(deviceId, RT_UTIL_TYPE_AICPU);
    // Currently, memory is not supported
    utilizationInfo->memoryUtilization = DEVICE_UTILIZATION_NOT_SUPPORT;
    ACL_LOG_INFO("successfully execute aclrtGetDeviceUtilizationRate, device is %d.", deviceId);
    return ACL_SUCCESS;
};

aclError aclrtGetDeviceCount(uint32_t *count)
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtGetDeviceCount);
    ACL_LOG_INFO("start to execute aclrtGetDeviceCount");
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(count);

    const rtError_t rtErr = rtGetDeviceCount(reinterpret_cast<int32_t *>(count));
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("get device count failed, runtime result = %d.", static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_INFO("successfully execute aclrtGetDeviceCount, get device count is %u.", *count);
    return ACL_SUCCESS;
}

aclError aclrtGetDeviceSatMode(aclrtFloatOverflowMode *mode)
{
    ACL_LOG_INFO("start to execute aclrtGetDeviceSatMode");
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(mode);
    rtFloatOverflowMode_t rtMode = RT_OVERFLOW_MODE_UNDEF;
    const rtError_t rtErr = rtGetDeviceSatMode(&rtMode);
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("rtGetDeviceSatMode failed, runtime result = %d.", static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    *mode = static_cast<aclrtFloatOverflowMode>(rtMode);
    ACL_LOG_INFO("successfully execute aclrtGetDeviceSatMode, mode is %d.", *mode);
    return ACL_SUCCESS;
}

aclError aclrtSetDeviceSatMode(aclrtFloatOverflowMode mode)
{
    ACL_LOG_INFO("start to execute aclrtSetDeviceSatMode, mode is %d", mode);
    const rtError_t rtErr = rtSetDeviceSatMode(static_cast<rtFloatOverflowMode_t>(mode));
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("rtSetDeviceSatMode failed, runtime result = %d.", static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_INFO("successfully execute aclrtSetDeviceSatMode, mode is %d", mode);
    return ACL_SUCCESS;
}

aclError aclrtGetOverflowStatus(void *outputAddr, size_t outputSize, aclrtStream stream)
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtGetOverflowStatus);
    ACL_LOG_INFO("start to execute aclrtGetOverflowStatus, outputSize = %lu", outputSize);
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(outputAddr);
    const rtError_t rtErr = rtGetDeviceSatStatus(outputAddr, outputSize, static_cast<rtStream_t>(stream));
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("rtGetDeviceSatStatus failed, runtime result = %d.", static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_INFO("successfully execute aclrtGetOverflowStatus");
    return ACL_SUCCESS;
}

aclError aclrtResetOverflowStatus(aclrtStream stream)
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtResetOverflowStatus);
    ACL_LOG_INFO("start to execute aclrtResetOverflowStatus");
    const rtError_t rtErr = rtCleanDeviceSatStatus(static_cast<rtStream_t>(stream));
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("rtCleanDeviceSatStatus failed, runtime result = %d.", static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    ACL_LOG_INFO("successfully execute aclrtResetOverflowStatus");
    return ACL_SUCCESS;
}

ACL_FUNC_VISIBILITY aclError aclrtQueryDeviceStatus(int32_t deviceId, aclrtDeviceStatus *deviceStatus)
{
    ACL_PROFILING_REG(acl::AclrtQueryDeviceStatus);
    ACL_LOG_INFO("start to execute aclrtQueryDeviceStatus with device id:%d", deviceId);
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(deviceStatus);
    rtDeviceStatus rtDevStatus = RT_DEVICE_STATUS_END;
    const rtError_t rtErr = rtDeviceStatusQuery(static_cast<uint32_t>(deviceId), &rtDevStatus);
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_WARN("rtDeviceStatusQuery failed, runtime result = %d.", static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    *deviceStatus = static_cast<aclrtDeviceStatus>(rtDevStatus);
    ACL_LOG_INFO("successfully execute aclrtQueryDeviceStatus");
    return ACL_SUCCESS;
}

aclError aclrtDeviceTaskAbort(int32_t deviceId, uint32_t timeout)
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtDeviceTaskAbort);
    ACL_LOG_INFO("start to execute aclrtDeviceTaskAbort on device %d, timeout %ums", deviceId, timeout);
    const rtError_t rtErr = rtDeviceTaskAbort(deviceId, timeout);
    if (rtErr != RT_ERROR_NONE) {
        ACL_LOG_ERROR("rtDeviceTaskAbort for device %d, failed, runtime result = %d.",
            deviceId, static_cast<int32_t>(rtErr));
        return ACL_GET_ERRCODE_RTS(rtErr);
    }
    return ACL_SUCCESS;
}
